
package cn.jbolt.ai.resource.vdb.service.impl;

import cn.hutool.core.util.StrUtil;
import cn.jbolt.ai.resource.llm.vo.CheckResourceVO;
import com.jboltai.resource.embedding.EmbeddingModelType;
import cn.jbolt.ai.resource.llm.vo.ResourceSelectVO;
import cn.jbolt.ai.resource.vdb.VDBType;
import cn.jbolt.util.Pager;
import com.baomidou.mybatisplus.core.conditions.query.QueryWrapper;
import com.baomidou.mybatisplus.core.conditions.update.UpdateWrapper;
import com.baomidou.mybatisplus.extension.plugins.pagination.Page;
import com.baomidou.mybatisplus.extension.service.impl.ServiceImpl;
import cn.jbolt.ai.resource.embedding.service.EmbeddingResourceService;
import cn.jbolt.ai.klb.entity.AiKnowledgeBase;
import cn.jbolt.ai.klb.service.AiKnowledgeBaseService;
import cn.jbolt.ai.resource.vdb.dao.VDBResourceMapper;
import cn.jbolt.ai.resource.vdb.entity.VDBResource;
import cn.jbolt.ai.resource.vdb.service.VDBResourceService;
import cn.jbolt.util.Result;
import com.jboltai.capability.vdb.JBoltVDB;
import com.jboltai.capability.vdb.VDBOptResult;
import com.jboltai.capability.vdb.types.milvus.MilvusVDBResource;
import com.jboltai.capability.vdb.types.tencent.TencentVDBResource;
import com.jboltai.resource.ResourceCenter;
import com.jboltai.resource.embedding.EmbeddingModel;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Transactional;

import java.util.ArrayList;
import java.util.Date;
import java.util.List;

@Service
public class VDBResourceServiceImpl extends ServiceImpl<VDBResourceMapper, VDBResource> implements VDBResourceService {
    @Autowired
    private AiKnowledgeBaseService knowledgeBaseService;
    @Autowired
    private EmbeddingResourceService embeddingResourceService;

    private static final Logger logger = LoggerFactory.getLogger(VDBResourceServiceImpl.class);

    @Override
    public Result<VDBResource> getResourceById(Long id) {
        VDBResource resource = this.getById(id);
        if (resource != null) {
            return Result.success(resource);
        } else {
            return Result.error("资源不存在");
        }
    }

    @Override
    public Result listByPage(Integer pageNumber, Integer pageSize, String keywords) {
        QueryWrapper<VDBResource> queryWrapper = new QueryWrapper<>();
        if (StrUtil.isNotBlank(keywords)) {
            queryWrapper.and(wrapper -> wrapper.like("name", keywords).
                    or().
                    like("type", keywords));
        }
        queryWrapper.orderByDesc("id");
        return Result.success(Pager.from(this.page(new Page<>(pageNumber, pageSize), queryWrapper)));
    }

    @Override
    @Transactional
    public Result add(VDBResource resource) {
        resource.setTimeout(60);//默认设置一个60s
        resource.setEnable(true);
        if (this.save(resource)) {
            Result result = registerResource(resource);
            if (result.isError()) {
//                throw new RuntimeException("资源注册失败");
                return Result.error("资源注册失败");
            }
            return Result.success("添加成功");
        } else {
            return Result.error("添加失败");
        }
    }

    @Override
    public Result check(VDBResource resource) {
        CheckResourceVO checkResourceVO = new CheckResourceVO();
        // 1. 注册资源
        Long resourceId = System.currentTimeMillis();
        resource.setId(resourceId);
        resource.setTimeout(60);//默认设置一个60s
        Result result = registerResource(resource);
        if (result.isError()) {
            checkResourceVO.setResult(false);
            checkResourceVO.setMsg(result.getMsg());
        }
        try {
            //2.检测数据库
            VDBOptResult<Boolean> booleanVDBOptResult = JBoltVDB.use(resourceId.toString()).checkDatabaseExist(resource.getDatabaseName());
            if (booleanVDBOptResult.isSuccess()) {
                checkResourceVO.setResult(true);
            } else {
                checkResourceVO.setResult(false);
                checkResourceVO.setMsg(booleanVDBOptResult.getFullMsg());
            }
        } catch (Exception e) {
            logger.error("检测VDB失败:", e);
            checkResourceVO.setResult(false);
            checkResourceVO.setMsg(e.getMessage());
        } finally {
            // 2. 删除资源
            unRegisterResource(resource);
            return Result.success(checkResourceVO);
        }

    }

    @Override
    @Transactional
    public Result update(VDBResource resource) {
        if (updateById(resource)) {
            if (resource.getEnable()) {
                unRegisterResource(resource);
                Result result = registerResource(resource);
                if (result.isError()) {
//                    throw new RuntimeException("资源注册失败");
                    return Result.error("资源注册失败");
                }
            }
            return Result.success();
        } else {
            return Result.error("更新失败");
        }
    }

    @Override
    @Transactional
    public Result del(Long id) {
        Result result = knowledgeBaseService.getListByVDBResourceId(id);
        if (result.isSuccess()) {
            List<AiKnowledgeBase> knowledgeBaseList = (List<AiKnowledgeBase>) result.getData();
            if (knowledgeBaseList.size() != 0) {
                return Result.error("该资源正在被使用，无法删除");
            }
        }
        Result<VDBResource> VDBResult = getResourceById(id);
        if (VDBResult.isError()) {
            return VDBResult;
        }
        VDBResource resource = VDBResult.getData();
        if (resource.getEnable()) {
            unRegisterResource(resource);
        }
        if (removeById(id)) {
            return Result.success();
        } else {
            return Result.error("删除失败");
        }
    }

    @Override
    @Transactional
    public Result toggleEnable(Long id) {
        VDBResource resource = this.getById(id);
        Boolean enable = !resource.getEnable();
        if (this.update(new UpdateWrapper<VDBResource>().eq("id", id).set("enable", enable))) {
            if (enable) {
                Result result = registerResource(resource);
                if (result.isError()) {
//                    throw new RuntimeException("资源注册失败");
                    return Result.error("资源注册失败");
                }
            } else {
                unRegisterResource(resource);
            }
            return Result.success();
        } else {
            return Result.error("更新启用状态失败");
        }
    }


    /**
     * 注册资源
     *
     * @param newResource
     */
    public Result registerResource( VDBResource newResource) {
        logger.info("VDB资源--{}--注册资源", newResource.getType());
        try {
            //判断是否是ollama
            // 使用 indexOf() 找到第一个冒号的位置
            String[] embedding = newResource.getEmbedding().split(":", 2);
            // 使用 substring() 方法来获取两个部分
            String modelType = embedding[0];
            String defaultModel = embedding.length == 2 ? embedding[1] : null;
            if (StrUtil.isEmpty(modelType)) {
                return Result.error("请先选择资源");
            }
            Result embeddingResult;
            if (modelType.equals(EmbeddingModelType.OLLAMA.getSn()) || modelType.equals(EmbeddingModelType.SILICONFLOW.getSn())) {
                Integer dimension = embeddingResourceService.getResourceByTypeAndDefaultModel(modelType, defaultModel).getDimension();
                embeddingResult = embeddingResourceService.getOrCreateModel(modelType,
                        defaultModel,
                        dimension);
            } else {
                embeddingResult = embeddingResourceService.getOrCreateModel(modelType,
                        null, null);
            }
            if (embeddingResult.isError()) {
                return embeddingResult;
            }
            EmbeddingModel embeddingModel = (EmbeddingModel) embeddingResult.getData();

            //判断选用的vdb类型
            if (newResource.getType().equals(cn.jbolt.ai.resource.vdb.VDBType.MILVUS.getSn())) {
                //先实例化MilvusVDBResource
                MilvusVDBResource milvusVDBResource =
                        new MilvusVDBResource(newResource.getId().toString(), newResource.getUrl(),
                                newResource.getUsername(), newResource.getPassword(), embeddingModel);
                milvusVDBResource.setDatabase(newResource.getDatabaseName());
                ResourceCenter.registerVDB(milvusVDBResource);
            } else if (newResource.getType().equals(cn.jbolt.ai.resource.vdb.VDBType.TENCENT.getSn())) {
                TencentVDBResource tencentVDBResource = new TencentVDBResource(newResource.getId().toString(), newResource.getUrl(), newResource.getUsername(), newResource.getPassword(), embeddingModel);
                tencentVDBResource.setDatabase(newResource.getDatabaseName());
                ResourceCenter.registerVDB(tencentVDBResource);
            }

            //尝试将数据库建起来
            try {
                VDBOptResult<Boolean> vdbOptResult = JBoltVDB.use(newResource.getId().toString()).checkDatabaseExist(newResource.getDatabaseName());
                if (vdbOptResult.isSuccess() && vdbOptResult.getResult() == false) {
                    JBoltVDB.use(newResource.getId().toString()).createDatabase(newResource.getDatabaseName());
                }
            } catch (Exception e) {
                logger.error("VDB资源--{}--创建数据库失败,异常原因:{}", newResource.getType(), e.getMessage());
            }

        } catch (Exception e) {
            String message = e.getMessage();
            logger.error("VDB资源--{}--注册资源失败,异常原因:{}", newResource.getType(), message);
            return Result.error(message);
        }
        return Result.success();

    }

    /**
     * 移除资源
     *
     * @param resource 资源
     */
    public void unRegisterResource(VDBResource resource) {

        logger.info("VDB资源--{}--移除资源", resource.getType());
        try {
            ResourceCenter.unregisterVDB(resource.getId().toString());
        } catch (Exception e) {
            logger.error("VDB资源--{}--移除资源失败,异常原因:{}", resource.getType(), e.getMessage());
            throw e;
        }
    }

    /**
     * 注册所有资源,是在项目启动时候调用
     */
    public void registerAllResource() {
        //只启动开启状态的资源
        this.list(new QueryWrapper<VDBResource>().eq("enable", true)).
                forEach(this::registerResource);
    }

    @Override
    public Result listAllEnabledType() {
        return Result.success(convertToTypeOptions(this.list(new QueryWrapper<VDBResource>()
                .eq("enable", true).orderByDesc("id"))));
    }

    /**
     * 根据类型转换成下拉选列表
     *
     * @param list
     * @return
     */
    public List<ResourceSelectVO> convertToTypeOptions(List<VDBResource> list) {
        List<ResourceSelectVO> selectList = new ArrayList<>();
        for (VDBResource resource : list) {
            ResourceSelectVO selectVO = new ResourceSelectVO();
            selectVO.setSn(resource.getType());
            selectVO.setName(VDBType.getBySn(resource.getType()).getName() + " : " + (StrUtil.isNotBlank(resource.getName()) ? resource.getName() : "默认"));
            selectVO.setId(resource.getId());
            selectList.add(selectVO);
        }
        return selectList;
    }
}
